How to extend a canned TensorFlow Estimator

Or how to add more evaluation metrics and pass through instance keys when using a canned estimator

Prebuilt (“canned”) estimators like DNNRegressor and DNNLinearCombinedRegressor make life easy when writing TensorFlow programs. For example, here is a complete TensorFlow program to train a TensorFlow model starting from an in-memory Pandas DataFrame:

import tensorflow as tf
featcols = [
  tf.feature_column.numeric_column("sq_footage")
]
model = tf.estimator.LinearRegressor(featcols, './model_trained')
def train_input_fn(df, num_epochs):  # a Pandas dataframe
  return tf.estimator.inputs.pandas_input_fn(
    x = df,
    y = df['price'],
    num_epochs=num_epochs)
model.train(train_input_fn(df,100), steps=1000)

If you want to use TensorFlow like scikit-learn, the number of lines of code is pretty much the same.

Robust Estimators

In reality, though, you want to be more robust than this. You want to read sharded files that don’t all fit into memory, batch them, etc. This involves a bit more code in your input function using the Dataset API in TensorFlow. And you want to checkpoint and evaluate every once in a while as you train. So, you don’t want just call train, you want to call train_and_evaluate.

Prebuilt Estimators are pretty straightforward, but also quite opinionated — they choose a specific way of doing things and give you a bunch of stuff out of the box. If you want more flexibility? Well, write a custom estimator (see example)!

But there is a gray area in between. You can use a canned estimator, and still get some amount of flexibility. How? Use the capabilities in extenders.py. In this blog post, I’ll show you two such capabilities (incidentally, these are two very frequently asked questions).

You can extend a canned estimator too!

How to add an extra evaluation metric

When you use a DNNRegressor, the metric that is part of the evaluation loop is only average_loss, which in this case happens to be the RMSE. But what if you want more metrics? Simply wrap the estimator by add_metrics as follows:

def my_rmse(labels, predictions):
    pred_values = predictions['predictions']
    return {'rmse': tf.metrics.root_mean_squared_error(labels, pred_values)}
def train_and_evaluate(output_dir):
    
    estimator = tf.estimator.DNNLinearCombinedRegressor(
        ...)
    
    estimator = tf.contrib.estimator.add_metrics(estimator, my_rmse)  
    ...
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

Let’s unpack that a bit:

1. Write a function that computes the metric(s) of interest from the labels and the predicted values and returns a dictionary.
2. Pass that function along with your original estimator to add_metrics
3. Use the returned estimator from that point onwards

Note that the my_rmse method above receives the labels and the entire predictions tensor. This way, if you are doing classification and some of your metrics need the predicted class and others need the predicted probabilities, you can compute what you want.

See a complete example here.

How to forward keys and input features to the output

The predictions dict returned by a canned Estimator has only the predictions in it (duh!). But there are lots of situations when you want the predictions to contain a bit of the input also. For example, you may have a unique key associated with each row, and you want that key to be part of the predictions so that you know which row the prediction corresponds to. Or you might want to compute evaluations based on specific input values (e.g. compute RMSE of baby weight based on whether the baby is pre-term or full-term).

To do that, wrap the estimator by forward_features as follows:

estimator = tf.contrib.estimator.forward_features(estimator, KEY_COLUMN)

You can pass in either a single string (as I did above) or a list of strings. Ideally, that’s all you need. However, there are two caveats.

The first caveat is that the key column above gets added to the output of predictions, but not to the export signature, so it will work while the model is in memory but not if it is restored after training (this makes no sense: I have submitted a pull request to the TensorFlow repo to fix this). Until then, you need this workaround:

def forward_key_to_export(estimator):
    estimator = tf.contrib.estimator.forward_features(estimator, KEY_COLUMN)
    # return estimator
## This shouldn't be necessary (I've filed CL/187793590 to update extenders.py with this code)
    config = estimator.config
    def model_fn2(features, labels, mode):
      estimatorSpec = estimator._call_model_fn(features, labels, mode, config=config)
      if estimatorSpec.export_outputs:
        for ekey in ['predict', 'serving_default']:
          estimatorSpec.export_outputs[ekey] = \
            tf.estimator.export.PredictOutput(estimatorSpec.predictions)
      return estimatorSpec
    return tf.estimator.Estimator(model_fn=model_fn2, config=config)
    ##

The second caveat is that you can not both pass-through a feature *and* use it as input to your model. Either consume the input or pass it through. The workaround here is to duplicate such columns that you want in your input functions using tf.identity:

features['gw'] = tf.identity(features['gestation_weeks'])

Hope this helps!